import torch
import torch.nn as nn
from torch_geometric.nn import GATConv


class M_GAT(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.attn1 = GATConv(in_dim, out_dim)
        # self.attn2 = GAT()

    def forward(self, x, edge_index):
        out = self.attn1(x, edge_index)
        return out


if __name__ == "__main__":
    h = torch.randn(4, 8)  # 4个节点，每个节点维度为8的特征
    edge_index = torch.tensor([[0, 1],
                               [1, 0]])
    # 对边矩阵(2,E) E为边的数量，第一行为边的起点，第二行为边的目标节点

    gat = M_GAT(8, 16)
    h1 = gat(h, edge_index)


    pass
